"""
测试训练管理器模块
"""
import os
import sys
import unittest
import numpy as np
import torch
import torch.nn as nn

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')))

from sichuanmajiang.model.training_manager import (
    TrainingManager, train_model, evaluate_model, export_trained_model,
    create_dummy_optimizer, create_dummy_scheduler
)
from sichuanmajiang.model.model import MahjongModel, create_mahjong_model
from sichuanmajiang.data.data_loader import MahjongDataset
from sichuanmajiang.tests.test_utils import (
    get_test_device, create_test_directory, remove_test_directory,
    create_test_features
)


class TestTrainingManager(unittest.TestCase):
    """
    测试TrainingManager类
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.device = get_test_device()
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_training')
        create_test_directory(self.test_dir)
        
        # 模型配置
        self.input_dim = 324
        self.action_dim = 6
        
        # 创建测试数据
        self.num_train_samples = 100
        self.num_val_samples = 20
        
        # 生成训练数据
        self.train_features = np.random.rand(self.num_train_samples, self.input_dim).astype(np.float32)
        self.train_actions = np.random.randint(0, self.action_dim, self.num_train_samples)
        self.train_rewards = np.random.rand(self.num_train_samples).astype(np.float32)
        
        # 生成验证数据
        self.val_features = np.random.rand(self.num_val_samples, self.input_dim).astype(np.float32)
        self.val_actions = np.random.randint(0, self.action_dim, self.num_val_samples)
        self.val_rewards = np.random.rand(self.num_val_samples).astype(np.float32)
        
        # 创建数据集
        self.train_dataset = MahjongDataset(
            features=self.train_features,
            actions=self.train_actions,
            rewards=self.train_rewards
        )
        self.val_dataset = MahjongDataset(
            features=self.val_features,
            actions=self.val_actions,
            rewards=self.val_rewards
        )
        
        # 创建训练管理器
        self.manager = TrainingManager(
            model=None,  # 稍后设置
            train_dataset=self.train_dataset,
            val_dataset=self.val_dataset,
            device=self.device,
            log_dir=self.test_dir
        )
        
        # 创建模型
        self.model = create_mahjong_model(
            input_dim=self.input_dim,
            action_dim=self.action_dim,
            hidden_dim=32,  # 使用较小的维度进行测试
            num_res_blocks=1
        )
        self.model.to(self.device)
        self.manager.model = self.model
        
        # 配置训练参数
        self.manager.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
        self.manager.scheduler = torch.optim.lr_scheduler.StepLR(self.manager.optimizer, step_size=10, gamma=0.1)
        self.manager.policy_loss_fn = nn.CrossEntropyLoss()
        self.manager.value_loss_fn = nn.MSELoss()
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_setup_logging(self):
        """
        测试日志设置
        """
        # 设置日志
        self.manager.setup_logging()
        
        # 检查日志目录是否创建
        self.assertTrue(os.path.exists(os.path.join(self.test_dir, 'tensorboard')))
    
    def test_build_model(self):
        """
        测试模型构建
        """
        # 构建新模型
        new_model = self.manager.build_model(
            model_type='mahjong_model',
            input_dim=self.input_dim,
            action_dim=self.action_dim,
            hidden_dim=32,
            num_res_blocks=1
        )
        
        # 检查模型类型
        self.assertIsInstance(new_model, MahjongModel)
    
    def test_setup_optimizer(self):
        """
        测试优化器设置
        """
        # 设置优化器
        optimizer = self.manager.setup_optimizer(optimizer_type='adam', lr=0.001)
        
        # 检查优化器类型
        self.assertIsInstance(optimizer, torch.optim.Adam)
    
    def test_setup_scheduler(self):
        """
        测试学习率调度器设置
        """
        # 设置优化器
        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
        
        # 设置调度器
        scheduler = self.manager.setup_scheduler(optimizer, scheduler_type='step', step_size=10, gamma=0.1)
        
        # 检查调度器类型
        self.assertIsInstance(scheduler, torch.optim.lr_scheduler.StepLR)
    
    def test_setup_loss_functions(self):
        """
        测试损失函数设置
        """
        # 设置损失函数
        policy_loss_fn, value_loss_fn = self.manager.setup_loss_functions(
            policy_loss_type='cross_entropy',
            value_loss_type='mse'
        )
        
        # 检查损失函数类型
        self.assertIsInstance(policy_loss_fn, nn.CrossEntropyLoss)
        self.assertIsInstance(value_loss_fn, nn.MSELoss)
    
    def test_train_one_epoch(self):
        """
        测试单轮训练
        """
        # 训练一轮
        epoch_loss, epoch_policy_loss, epoch_value_loss = self.manager.train_one_epoch(
            batch_size=16,
            policy_weight=0.5,
            value_weight=0.5
        )
        
        # 检查损失是否为有效数值
        self.assertIsInstance(epoch_loss, float)
        self.assertIsInstance(epoch_policy_loss, float)
        self.assertIsInstance(epoch_value_loss, float)
        
        # 检查损失是否合理
        self.assertGreaterEqual(epoch_loss, 0)
        self.assertGreaterEqual(epoch_policy_loss, 0)
        self.assertGreaterEqual(epoch_value_loss, 0)
    
    def test_evaluate(self):
        """
        测试评估功能
        """
        # 评估模型
        metrics = self.manager.evaluate(batch_size=16)
        
        # 检查评估指标
        self.assertIn('loss', metrics)
        self.assertIn('policy_loss', metrics)
        self.assertIn('value_loss', metrics)
        self.assertIn('accuracy', metrics)
        
        # 检查指标是否为有效数值
        self.assertIsInstance(metrics['loss'], float)
        self.assertIsInstance(metrics['policy_loss'], float)
        self.assertIsInstance(metrics['value_loss'], float)
        self.assertIsInstance(metrics['accuracy'], float)
        
        # 检查准确性是否在合理范围内
        self.assertGreaterEqual(metrics['accuracy'], 0)
        self.assertLessEqual(metrics['accuracy'], 1)
    
    def test_train(self):
        """
        测试训练功能（简短训练）
        """
        # 进行简短训练（1个epoch）
        history = self.manager.train(
            num_epochs=1,
            batch_size=16,
            policy_weight=0.5,
            value_weight=0.5,
            save_best=True,
            save_dir=self.test_dir
        )
        
        # 检查训练历史
        self.assertIsInstance(history, dict)
        self.assertIn('train_losses', history)
        self.assertIn('val_losses', history)
        self.assertIn('train_accuracies', history)
        self.assertIn('val_accuracies', history)
        
        # 检查是否保存了最佳模型
        self.assertTrue(os.path.exists(os.path.join(self.test_dir, 'best_model.pt')))
    
    def test_test(self):
        """
        测试测试功能
        """
        # 创建测试数据
        test_features = np.random.rand(20, self.input_dim).astype(np.float32)
        test_actions = np.random.randint(0, self.action_dim, 20)
        test_rewards = np.random.rand(20).astype(np.float32)
        
        test_dataset = MahjongDataset(
            features=test_features,
            actions=test_actions,
            rewards=test_rewards
        )
        
        # 测试模型
        test_metrics = self.manager.test(test_dataset, batch_size=16)
        
        # 检查测试指标
        self.assertIn('loss', test_metrics)
        self.assertIn('policy_loss', test_metrics)
        self.assertIn('value_loss', test_metrics)
        self.assertIn('accuracy', test_metrics)
    
    def test_export_model(self):
        """
        测试导出模型功能
        """
        # 导出模型
        export_path = os.path.join(self.test_dir, 'exported_model.onnx')
        self.manager.export_model(export_path)
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(export_path))


class TestHelperFunctions(unittest.TestCase):
    """
    测试辅助函数
    """
    def setUp(self):
        """
        测试前的设置
        """
        self.device = get_test_device()
        self.test_dir = os.path.join(os.path.dirname(__file__), 'test_functions')
        create_test_directory(self.test_dir)
        
        # 模型配置
        self.input_dim = 324
        self.action_dim = 6
        
        # 创建模型
        self.model = create_mahjong_model(
            input_dim=self.input_dim,
            action_dim=self.action_dim,
            hidden_dim=32,
            num_res_blocks=1
        )
        self.model.to(self.device)
        
        # 创建测试数据
        self.num_train_samples = 50
        self.num_val_samples = 10
        self.num_test_samples = 10
        
        # 生成数据
        self.train_features = np.random.rand(self.num_train_samples, self.input_dim).astype(np.float32)
        self.train_actions = np.random.randint(0, self.action_dim, self.num_train_samples)
        self.train_rewards = np.random.rand(self.num_train_samples).astype(np.float32)
        
        self.val_features = np.random.rand(self.num_val_samples, self.input_dim).astype(np.float32)
        self.val_actions = np.random.randint(0, self.action_dim, self.num_val_samples)
        self.val_rewards = np.random.rand(self.num_val_samples).astype(np.float32)
        
        self.test_features = np.random.rand(self.num_test_samples, self.input_dim).astype(np.float32)
        self.test_actions = np.random.randint(0, self.action_dim, self.num_test_samples)
        self.test_rewards = np.random.rand(self.num_test_samples).astype(np.float32)
        
        # 创建数据集
        self.train_dataset = MahjongDataset(
            features=self.train_features,
            actions=self.train_actions,
            rewards=self.train_rewards
        )
        self.val_dataset = MahjongDataset(
            features=self.val_features,
            actions=self.val_actions,
            rewards=self.val_rewards
        )
        self.test_dataset = MahjongDataset(
            features=self.test_features,
            actions=self.test_actions,
            rewards=self.test_rewards
        )
    
    def tearDown(self):
        """
        测试后的清理
        """
        # 移除测试目录
        remove_test_directory(self.test_dir)
    
    def test_train_model(self):
        """
        测试train_model函数
        """
        # 训练模型
        history = train_model(
            model=self.model,
            train_dataset=self.train_dataset,
            val_dataset=self.val_dataset,
            num_epochs=1,
            batch_size=16,
            lr=0.001,
            device=self.device,
            log_dir=self.test_dir
        )
        
        # 检查训练历史
        self.assertIsInstance(history, dict)
        self.assertIn('train_losses', history)
        self.assertIn('val_losses', history)
    
    def test_evaluate_model(self):
        """
        测试evaluate_model函数
        """
        # 评估模型
        metrics = evaluate_model(
            model=self.model,
            dataset=self.val_dataset,
            batch_size=16,
            device=self.device
        )
        
        # 检查评估指标
        self.assertIn('loss', metrics)
        self.assertIn('accuracy', metrics)
    
    def test_export_trained_model(self):
        """
        测试export_trained_model函数
        """
        # 导出模型
        export_path = os.path.join(self.test_dir, 'function_exported_model.onnx')
        export_trained_model(
            model=self.model,
            export_path=export_path,
            input_dim=self.input_dim
        )
        
        # 检查文件是否存在
        self.assertTrue(os.path.exists(export_path))
    
    def test_create_dummy_optimizer(self):
        """
        测试create_dummy_optimizer函数
        """
        # 创建优化器
        optimizer = create_dummy_optimizer(self.model)
        
        # 检查类型
        self.assertIsInstance(optimizer, torch.optim.Optimizer)
    
    def test_create_dummy_scheduler(self):
        """
        测试create_dummy_scheduler函数
        """
        # 创建优化器
        optimizer = torch.optim.Adam(self.model.parameters())
        
        # 创建调度器
        scheduler = create_dummy_scheduler(optimizer)
        
        # 检查类型
        self.assertIsInstance(scheduler, torch.optim.lr_scheduler._LRScheduler)


if __name__ == '__main__':
    unittest.main()
